#!/usr/bin/env python3
# -*- coding:utf-8 -*-

import os
import sys
import argparse
import re
import enum
import traceback
import time

import csv

import logging
log_format = "%(levelname)s: %(filename)s-%(lineno)d/%(funcName)s: %(message)s"
log_format = "%(levelname)s: %(lineno)d/%(funcName)s: %(message)s"

logging.basicConfig(format=log_format)
log = logging.getLogger("[logdebug]")
log.setLevel(logging.DEBUG)

@enum.unique
class RowType(enum.Enum):
    ERROR       = 'error'
    FILE        = 'file'
    FUNCTION    = 'func'
    FUNC_DECL   = 'decl'
    FUNC_CALL   = 'call'


class CIParserBase():

    @staticmethod
    def preprocess_line(line, prefix):
        line = line.strip(prefix).strip()
        line = line.strip(':').strip()
        line = line.strip('{').strip('}').strip()
        return line

    @staticmethod
    def strip_symbol(sym):
        ext = False
        post = sym.split(':')[-1]
        # gcc optimize
        if '.' in post:
            ext = True
            post = post.split('.')[0]
        if ':' in sym:
            head = sym.split(':')[0]
            return '{}:{}'.format(head, post), ext
        return post, ext

    @staticmethod
    def row_replace(row, prefix):
        if not prefix:
            return tuple(row)
        tmp = []
        for x in row:
            if isinstance(x, str):
                x = x.replace(prefix, '')
            tmp.append(x)
        return tuple(tmp)

    @staticmethod
    def __skip_line(line):
        if not line:
            return True
        if '}' == line:
            return True
        if '__indirect_call' in line:
            return True
        if '__stack_chk_fail' in line:
            return True
        return False

    @classmethod
    def __get_flc(cls, label):
        filepath = ''
        lineno, column = 0, 0
        if label.count(':') == 2:
            tmps = label.split(':')
            filepath = tmps[0].split(r'\n')[-1]
            lineno = int(tmps[1])
            column = int(tmps[2])
        return filepath, lineno, column

    @classmethod
    def process_line(cls, line):
        def __strip(s, st):
            return s.strip().strip(st).strip()
        '''
        edge: { sourcename: "kernel/module.c:resolve_symbol"
        targetname: "_printk"
        label: "kernel/module.c:1382:3" }
        '''
        row_type = RowType.ERROR
        ret = {}
        if '{' not in line:
            return row_type, ret

        slices = line.split('{')
        head = __strip(slices[0], ':')
        slices = __strip(slices[-1], '}')
        slices = slices.replace('" ', ':"')
        slices = slices.replace(' ', '')
        slices = slices.split(':"')

        key = None
        for s in slices:
            if not key:
                key = s.strip()
            else:
                s = s.strip('"')
                ret[key] = s
                key = None

        if head == 'graph':
            row_type = RowType.FILE
        elif head == 'node':
            if 'shape' in line:
                row_type = RowType.FUNC_DECL
            else:
                row_type = RowType.FUNCTION
        elif head == 'edge':
            row_type = RowType.FUNC_CALL

        return row_type, ret

    @classmethod
    def __parse(cls, filename, prefix=None):
        srcfiles, decls, funcs, links = set(), set(), set(), set()
        with open(filename, 'r', encoding='utf-8') as cifd:
            for idx, line in enumerate(cifd):
                try:
                    line = line.strip()
                    if cls.__skip_line(line):
                        continue

                    row_type, ret = cls.process_line(line)
                    if row_type == RowType.FILE:
                        filepath = ret['title']
                        row = ('file', filepath, 0, 0, '', '')
                        row = cls.row_replace(row, prefix)
                        srcfiles.add(row)

                    elif row_type == RowType.FUNC_DECL:
                        wdecl = ret['title']
                        label = ret['label']
                        if '<built-in>' in label:
                            label = ''

                        filepath, lineno, column = cls.__get_flc(label)
                        wdecl, ext = cls.strip_symbol(wdecl)
                        decl = wdecl.split(':')[-1]

                        row = ('decl', filepath, lineno, column, decl, wdecl, '', '')
                        row = cls.row_replace(row, prefix)
                        decls.add(row)

                    elif row_type == RowType.FUNCTION:
                        wfunc = ret['title']
                        label = ret['label']
                        if '<built-in>' in label:
                            label = ''
                        filepath, lineno, column = cls.__get_flc(label)
                        wfunc, ext = cls.strip_symbol(wfunc)
                        func = wfunc.split(':')[-1]

                        row = ('func', filepath, lineno, column, func, wfunc, '', '')
                        row = cls.row_replace(row, prefix)
                        funcs.add(row)

                    elif row_type == RowType.FUNC_CALL:
                        wsrc = ret['sourcename']
                        wdst = ret['targetname']

                        wsrc, ext = cls.strip_symbol(wsrc)
                        wdst, ext = cls.strip_symbol(wdst)
                        src, dst = wsrc.split(':')[-1], wdst.split(':')[-1]

                        label = ''
                        if 'label' in ret.keys():
                            label = ret['label']
                        filepath, lineno, column = cls.__get_flc(label)

                        row = ('call', filepath, lineno, column, src, wsrc, dst, wdst)
                        row = cls.row_replace(row, prefix)
                        links.add(row)

                    else:
                        log.warning('file: {}, line: {}'.format(filename, idx + 1))
                except Exception as e:
                    log.warning(e)
                    log.warning(line)
                    log.debug(traceback.format_exc())
        return srcfiles, decls, funcs, links

    @classmethod
    def merge(cls, srcfiles, decls, funcs, links):
        # ('decl', filepath, lineno, column, decl, wdecl, '', '')
        # ('func', filepath, lineno, column, func, wfunc, '', '')

        wfuncnames = {row[-3] for row in funcs}
        funcnames = {row[-4] for row in funcs}

        n_decls = set()
        for row in decls:
            wdecl = row[-3]
            if wdecl not in wfuncnames:
                n_decls.add(row)

        decl_dict = {}
        for row in sorted(n_decls, reverse=True):
            wdecl = row[-3]
            if wdecl not in decl_dict:
                decl_dict[wdecl] = [row]
            else:
                filepath = row[1]
                if filepath:
                    decl_dict[wdecl].append(row)

        for wdecl in decl_dict:
            if len(decl_dict[wdecl]) == 1:
                continue
            # copy
            for row in decl_dict[wdecl][:]:
                filepath = row[1]
                if not filepath:
                    decl_dict[wdecl].remove(row)

        n_decls = set()
        for rows in decl_dict.values():
            n_decls |= set(rows)

        return srcfiles, n_decls, funcs, links

    @classmethod
    def parse_s(cls, filename, prefix=None, merge=False):
        srcfiles, decls, funcs, links = cls.__parse(filename, prefix)
        if merge:
            return cls.merge(srcfiles, decls, funcs, links)
        return srcfiles, decls, funcs, links

    @classmethod
    def save_csv_s(cls, filename, rows, sortkey=None):
        with open(filename, 'w') as fd:
            header = ('type', 'file', 'line', 'column',
                      'symbol', 'wsymbol', 'callee', 'wcallee')
            if sortkey:
                index = header.index(sortkey)
                if index >= 0:
                    rows = sorted(rows, key=lambda x:x[index])
            wr = csv.writer(fd)
            wr.writerow(header)
            wr.writerows(rows)


class CIParser(CIParserBase):
    def __init__(self):
        self.srcfiles = set()
        self.decls = set()
        self.funcs = set()
        self.links = set()

    def rows(self):
        return self.decls | self.funcs | self.links

    def save_csv(self, filename, sortkey=None):
        self.save_csv_s(filename, self.rows(), sortkey)

    def parse(self, filename, prefix=None, merge=False):
        srcfiles, decls, funcs, links = self.parse_s(filename, prefix)
        self.srcfiles |= srcfiles
        self.decls |= decls
        self.funcs |= funcs
        self.links |= links
        if merge:
            self.srcfiles, self.decls, self.funcs, self.links = self.merge(self.srcfiles, self.decls, self.funcs, self.links)
        return self.srcfiles, self.decls, self.funcs, self.links

class MultiCIParser(CIParserBase):
    def __init__(self):
        self.srcfiles = set()
        self.decls = set()
        self.funcs = set()
        self.links = set()

    def rows(self):
        return self.decls | self.funcs | self.links

    def save_csv(self, filename, sortkey=None):
        self.save_csv_s(filename, self.rows(), sortkey)

    @classmethod
    def get_cilist(cls, filename):
        cilist = []
        with open(filename, 'r') as fd:
            for i, line in enumerate(fd):
                line = line.strip()
                if re.match(r'.*\.[0-9]*\.ci$', line):
                    continue
                cilist.append(line)
        return cilist

    @classmethod
    def parse_list_s(cls, cifiles, prefix=None, merge=False):
        srcfiles, decls, funcs, links = set(), set(), set(), set()
        for i, filename in enumerate(cifiles):
            if re.match(r'.*\.[0-9]*\.ci$', filename):
                continue
            try:
                _files, _decls, _funcs, _links = super().parse_s(filename, prefix, merge=False)
                srcfiles |= _files
                decls |= _decls
                funcs |= _funcs
                links |= _links
            except Exception as e:
                log.warning(e)
                log.warning(traceback.format_exc())
        if merge:
            return cls.merge(srcfiles, decls, funcs, links)
        return srcfiles, decls, funcs, links

    @classmethod
    def parse_cilist_s(cls, filename, prefix=None, merge=False):
        cifiles = cls.get_cilist(filename)
        return cls.parse_list_s(cifiles, prefix, merge)

    def parse_list(self, cifiles:list, prefix=None, merge=False):
        srcfiles, decls, funcs, links = self.parse_list_s(cifiles, prefix)
        self.srcfiles |= srcfiles
        self.decls |= decls
        self.funcs |= funcs
        self.links |= links
        if merge:
            self.srcfiles, self.decls, self.funcs, self.links = self.merge(self.srcfiles, self.decls, self.funcs, self.links)
        return self.srcfiles, self.decls, self.funcs, self.links

    def parse_cilist(self, filename, prefix=None, merge=False):
        cifiles = self.get_cilist(filename)
        return self.parse_list(cifiles, prefix, merge)

def main(args):
    cifiles = []

    if args.infmt == 'ci':
        cifiles = args.files
    elif args.infmt == 'list':
        for filename in args.files:
            cifiles += MultiCIParser.get_cilist(filename)

    t_start = time.time()
    cip = MultiCIParser()
    print('parsing {} file(s)'.format(len(cifiles)))
    cip.parse_list(cifiles, args.prefix, merge=True)
    rows = cip.rows()
    sortkey = args.sort
    print('parsing use {t_use:.2f}s'.format(t_use=(time.time() - t_start)))
    print('rows: {}'.format(len(rows)))
    cip.save_csv(args.output, sortkey)

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.description = '%(prog)s: CI files parser'

    parser.add_argument("-I", "--infmt", required=True,
                        help='Input formats(file suffix) are: ci/list')

    parser.add_argument("-o", "--output", required=True,
                        help="output file")

    parser.add_argument("-p", "--prefix",
                        help='path prefix, usually same as source code path')

    parser.add_argument("-s", "--sort", type=str,
                        help='sort by type,file,symbol,wsymbol,callee,wcallee')

    parser.add_argument("files", nargs="+", type=str,
                        help="input files")

    args = parser.parse_args()
    print(args)
    if args.infmt not in ('ci', 'list'):
        parser.print_help()
        os._exit(-1)
    main(args)
